# -*- coding: utf-8 -*-
"""
Created on 2019/11/22 下午5:03

@Project -> File: ode-neural-network -> eval_ode_net.py

@Author: luolei

@Describe: ODE Net模型评估
"""

import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import torch
import json

from lib import cols, col_bounds, discrete_t_steps
from lib.ode_net.train_ode_net import integrate
from mod.neural_network_model.nn import PartialDeriveNet


def load_models():
	"""载入已经训练好的模型"""
	with open('../../file/model/struc_params.json', 'r') as f:
		model_struc_params = json.load(f)
	
	model_path = '../../file/model/state_dict.pth'
	pretrained_model_dict = torch.load(model_path, map_location = 'cpu')
	
	input_size = model_struc_params['pd_net']['input_size']
	hidden_sizes = model_struc_params['pd_net']['hidden_sizes']
	output_size = model_struc_params['pd_net']['output_size']
	model = PartialDeriveNet(input_size, hidden_sizes, output_size)
	model.load_state_dict(pretrained_model_dict, strict = False)
	
	model.eval()
	return model


def vstack_arr(total, sub):
	"""竖向拼接数据"""
	
	if total is None:
		total = sub
	else:
		total = np.vstack((total, sub))
	
	return total


if __name__ == '__main__':
	# %% 载入模型文件
	pd_net = load_models()
	
	# %% 载入数据
	data = pd.read_csv('../../data/runtime/total_obs_data.csv')
	
	# 数据归一化
	for col in cols:
		bounds = col_bounds[col]
		data[col] = data[col].apply(lambda x: (x - bounds[0]) / (bounds[1] - bounds[0]))
	
	
	# %% 构造数据样本
	labels = list(data['label'].drop_duplicates())
	x0 = None
	for label in labels:
		sub_data = data[data.label == label].copy()
		
		sub_x0 = sub_data.iloc[: -1][cols].to_numpy()  # 获得初值x0
		
		sub_t0 = sub_data.iloc[: -1][['time']].to_numpy()  # 获得初始时间t0
		
		# 数据拼接
		x0 = vstack_arr(x0, sub_x0)
			
	x0 = torch.from_numpy(x0.astype(np.float32))
	
	# %% 模型预测
	t_init = data['time'].min()
	t_final = data['time'].max()
	dt = 0.01
	
	plt.figure(figsize = [5, 5])
	for i in range(60):
		x_init = x0[0: 1, :]
		
		if i != 0:
			x_rand = torch.cat(
				(
					torch.normal(0.06, std = torch.arange(0.2, 0.21, 0.2)),
					torch.normal(0.42, std = torch.arange(0.8, 0.81, 0.8))
				),
				dim = 0
			).reshape(1, -1)
			
			x_init = torch.add(x_init, x_rand)
		_, x_records = integrate(x_init, t_init, t_final, dt, pd_net)
		
		x1_pred_epoch = x_records.detach().cpu().numpy()
		
		if i == 0:
			c = '#ff7f0e'  # 默认橙黄色
		else:
			c = '#1f77b4'  # 默认蓝色
			
		plt.plot(x1_pred_epoch[:, 0, 0], x1_pred_epoch[:, 0, 1], c = c, linewidth = 0.6)
	plt.title('phase portrait')
	plt.savefig('../../graph/eval_effect.png', dpi = 450)


